package com.xjy.spring.thrift.core;

import com.xjy.spring.thrift.config.BaseServerConfig;
import org.I0Itec.zkclient.IZkChildListener;
import org.I0Itec.zkclient.ZkClient;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.apache.commons.lang3.reflect.ConstructorUtils;
import org.apache.commons.lang3.reflect.MethodUtils;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.springframework.aop.framework.ProxyFactory;
import org.springframework.beans.factory.FactoryBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.Assert;
import org.springframework.util.ClassUtils;
import org.springframework.util.CollectionUtils;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;

public class ThriftClientProxyFactoryBean<T> implements FactoryBean<T>, InitializingBean, MethodInterceptor {

    // thrift生成的service中Iface内部类的名称
    private static final String iFaceInnerClassName = "Iface";
    // thrift生成的service中Client内部类的名称
    private static final String clientInnerClassName = "Client";
    // thriftClientProxyFactoryBean产生的bean
    private T serviceProxy;
    // 连接到thrift的端口号
    private Integer port;
    // 连接到thrift的host
    private String host;
    // thrift生成的service的包名
    private String thriftServicePackageName;
    // 连接超时时间
    private Integer timeout = 3000;
    // 是否使用异步
    private boolean async;

    public boolean getAsync() {
        return async;
    }

    public void setAsync(boolean async) {
        this.async = async;
    }

    public String getThriftServicePackageName() {
        return thriftServicePackageName;
    }

    public void setThriftServicePackageName(String thriftServicePackageName) {
        this.thriftServicePackageName = thriftServicePackageName;
    }

    public String getHost() {
        return host;
    }

    public void setHost(String host) {
        this.host = host;
    }

    public Integer getPort() {
        return port;
    }

    public void setPort(Integer port) {
        this.port = port;
    }

    public Integer getTimeout() {
        return timeout;
    }

    public void setTimeout(Integer timeout) {
        this.timeout = timeout;
    }

    @Override
    public T getObject() throws Exception {
        return this.serviceProxy;
    }

    @Override
    public Class<?> getObjectType() {
        validate();
        Class<?> objectType;
        try {
            objectType = initClass();
        } catch (ClassNotFoundException e) {
            throw new RuntimeException("Can not find class " + thriftServicePackageName + "$" + clientInnerClassName, e);
        }
        return objectType;
    }

    @Override
    public boolean isSingleton() {
        return true;
    }

    // 定义一个 volatile 成员变量，用于保存最新的地址（考虑到该变量或许会被其它线程所修改，一旦修改后，该变量的值会影响到所有线程，所以保证最终可见性）
    private volatile Map<String, BaseServerConfig> thriftServerConfigMap = new HashMap<String, BaseServerConfig>();
    private ZkClient zkClient;
    private static final String CONNECT_STRING = "localhost:2181";
    private static final int SESSION_TIMEOUT = 5000;
    private static final int CONNECTION_TIMEOUT = 1000;
    private static final String ROOT_PATH = "/thrift";

    @Override
    public void afterPropertiesSet() throws Exception {
        validate();
        this.serviceProxy = initServiceProxy();
        // 连接zookeeper
        connectZookeeper();
        // 监听
        subscribeChildChanges();
        // 初始化thrift服务信息
        initThriftServiceInfo();
    }

    private T initServiceProxy() throws ClassNotFoundException {
        return (T)new ProxyFactory(org.apache.commons.lang3.ClassUtils.getClass(thriftServicePackageName + "$" + iFaceInnerClassName)
                , this).getProxy(ClassUtils.getDefaultClassLoader());
    }

    private void connectZookeeper() {
        if (zkClient == null) {
            zkClient = new ZkClient(CONNECT_STRING, SESSION_TIMEOUT, CONNECTION_TIMEOUT);
        }
    }

    private void subscribeChildChanges() {
        if (zkClient == null) {
            throw new RuntimeException("connect zookeeper error");
        }
        zkClient.subscribeChildChanges(ROOT_PATH, new IZkChildListener() {

            @Override
            public void handleChildChange(String parentPath, List<String> currentChildren) throws Exception {
                // 重新初始化
                if(CollectionUtils.isEmpty(currentChildren)) {
                    thriftServerConfigMap.clear();
                    return;
                }
                for (String node : currentChildren) {
                    // 获取子节点中的数据
                    BaseServerConfig config = zkClient.readData(ROOT_PATH + "/" + node);
                    thriftServerConfigMap.put(config.getHost() + ":" + config.getPort(), config);
                }
            }
        });
    }

    private void initThriftServiceInfo() {
        if (thriftServerConfigMap.isEmpty()) {
            List<String> currentChildren = zkClient.getChildren(ROOT_PATH);
            for (String node : currentChildren) {
                // 获取子节点中的数据
                BaseServerConfig config = zkClient.readData(ROOT_PATH + "/" + node);
                thriftServerConfigMap.put(config.getHost() + ":" + config.getPort(), config);
            }
        }
    }

    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
        Object result;
        TTransport transport = null;

        try {
            BaseServerConfig config = lookupServerConfig();
            if(config == null) {
                return null;
            }

            transport = openTTransport(config.getHost(), config.getPort());
            if(transport == null) {
                return null;
            }

            // 协议要和服务端一致
            TProtocol protocol = async ? new TCompactProtocol(new TFramedTransport(transport)) : new TBinaryProtocol(transport);

            Class<?> clientClass = initClass();

            result = invokeMethod(invocation, clientClass, protocol);
        } catch (TTransportException e) {
            throw new TTransportException("Open transport error. Host is " + host + " ,port is " + port, e);
        } catch (ClassNotFoundException e) {
            throw new ClassNotFoundException("Can not find class " + thriftServicePackageName + "$" + clientInnerClassName, e);
        } catch (InvocationTargetException e) {
            throw new InvocationTargetException(e, "Invoke method error.");
        } catch (NoSuchMethodException e) {
            throw new NoSuchMethodException("Invoke method error:\n" + e.getMessage());
        } catch (InstantiationException e) {
            throw new InstantiationException("Invoke method error:\n" + e.getMessage());
        } catch (IllegalAccessException e) {
            throw new IllegalAccessException("Invoke method error:\n" + e.getMessage());
        } catch (Throwable e) {
            if (e instanceof TException) {
                throw new TException("Thrift inner error", e);
            }
            throw e;
        } finally {
            closeTTransport(transport);
        }
        return result;
    }


    // 查找服务地址
    private BaseServerConfig lookupServerConfig() {
        if(thriftServerConfigMap.size() == 0) {
            System.err.println("No server found!");
            return null;
        }
        List<String> keys = new ArrayList<String>(thriftServerConfigMap.keySet());
        int size = keys.size();
        BaseServerConfig config;
        if (size == 1) {
            config = thriftServerConfigMap.get(keys.get(0));
        } else {
            // 随机算法
            config = thriftServerConfigMap.get(keys.get((ThreadLocalRandom.current().nextInt(size))));
        }
        return config;
    }


    private void validate() {
        Assert.hasLength(thriftServicePackageName, "thriftServicePackageName can not be null");
        Assert.hasLength(clientInnerClassName, "clientInnerClassName can not be null");
        Assert.hasLength(thriftServicePackageName, "thriftServicePackageName can not be null");
        Assert.notNull(port, "port can not be null");
        Assert.notNull(host, "host can not be null");
    }

    private TTransport initTTransport(String host, Integer port) {
        TTransport transport;
        if (timeout.equals(new Integer(-1))) {
            transport = new TSocket(host, port);
        } else {
            transport = new TSocket(host, port, timeout);
        }
        return transport;
    }

    private void closeTTransport(TTransport transport) {
        if (null != transport && transport.isOpen()) {
            transport.close();
        }
    }

    private TTransport openTTransport(String host, Integer port) throws TTransportException {
        TTransport transport = initTTransport(host, port);
        try {
            transport.open();
        } catch (TTransportException e) {
            e.printStackTrace();
            // 简单的重试机制 TODO add by xujieyang 优化
            BaseServerConfig config = lookupServerConfig();
            if(config == null) {
                return null;
            }
            openTTransport(config.getHost(), config.getPort());
        }
        return transport;
    }

    private Class<?> initClass() throws ClassNotFoundException {
        return org.apache.commons.lang3.ClassUtils.getClass(thriftServicePackageName + "$" + clientInnerClassName);
    }

    private Object invokeMethod(MethodInvocation invocation, Class<?> clientClass, TProtocol protocol)
            throws InvocationTargetException, NoSuchMethodException, InstantiationException, IllegalAccessException {
        Method method = invocation.getMethod();
        return MethodUtils.invokeMethod(ConstructorUtils.invokeConstructor(clientClass, protocol),
                method.getName(), invocation.getArguments());
    }

}
